/*
Copyright 2008, 2009 Joonas Lehtolahti

This file is part of Shared Datatypes and Utilities.

Shared Datatypes and Utilities is free software: you can redistribute it and/or modify
it under the terms of the GNU Lesser General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

Shared Datatypes and Utilities is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU Lesser General Public License for more details.

You should have received a copy of the GNU Lesser General Public License
along with Shared Datatypes and Utilities.  If not, see <http://www.gnu.org/licenses/>.
*/

#include "tree.h"

/* AVL Tree balancing function */
TreeNode* AVLTreeCheckNode(TreeNode*,BTree*);

/* AVL Tree balance counter */
int AVLTreeCountBalance(TreeNode*);

/* AVL Tree actual adding function */
TreeNode *BTreeAddActual(BTree *tree, void *x, TreeNode *p);


/**
 * newBTree
 *	Creates a new BTree pointer
 * 
 * Parameters:
 *	cmpfunc	a pointer to function used to compare items
 * 
 * Returns:
 *	the pointer to the new BTree structure
 *	NULL in case of error
 */
DLLEXPORT(BTree*) newBTree(CMPFUNC*cmpfunc) {
	BTree* x;
	if ((x = (BTree*)malloc(sizeof(BTree)))==NULL) return NULL;
	x->cmpfunc = cmpfunc;
	x->root = NULL;
	return x;
}

/**
 * BTreeFind
 *	Finds the specified element from a BTree
 * 
 * Parameters:
 *	tree	the BTree from which the element is to be searched
 *	x		the element to be searched for
 * 
 * Returns:
 *	a pointer to the TreeNode that contains the specific element
 *	NULL if the element was not found from the tree
 */
DLLEXPORT(TreeNode*) BTreeFind(BTree* tree, void* x) {
	if (tree && x) {
		TreeNode* n = tree->root;
		int (*comparator)(void*,void*) = tree->cmpfunc;
		int result;
		while (n) {
			result = comparator(x,n->x);
			if (result==0) return n;
			if (result<0) n=n->left;
			else n=n->right;
		}
	}
	return NULL;
}

/**
 * BTreeAdd
 *	Adds an element into BTree structure performing tree balancing
 *	using AVL methods when necessary. Fails if item exists.
 * 
 * Parameters:
 *	tree  the tree into which the element should be added to
 *	x     the pointer which is saved into the tree
 * 
 * Returns:
 *	the TreeNode of added element
 *  NULL if the element exists already
 */
DLLEXPORT(TreeNode*) BTreeAdd(BTree *tree, void *x) {
	if (tree) {
		/* Initialization */
		TreeNode* n = tree->root;
		TreeNode* p = NULL;
		int (*comparator)(void*,void*) = tree->cmpfunc;
		int cmpresult;
		/* Find correct place for the new node inside while loop */
		while (n) {
			p = n;
			cmpresult = comparator(x,n->x);
			if (cmpresult>0) {
				n = n->right;
			}
			else if (cmpresult<0){
				n = n->left;
			}
			else {
				return NULL;
			}
		}
		return BTreeAddActual(tree, x, p);
	}
	return NULL;
}

/**
 * BTreeFindOrAdd
 *	Adds an element into BTree structure performing tree balancing
 *	using AVL methods when necessary. If item exists, returns old node.
 * 
 * Parameters:
 *	tree  the tree into which the element should be added to
 *	x     the pointer which is saved into the tree
 * 
 * Returns:
 *	the TreeNode of found or added element
 *  NULL if the tree was empty or adding node failed (?)
 */
DLLEXPORT(TreeNode*) BTreeFindOrAdd(BTree *tree, void *x) {
	if (tree) {
		/* Initialization */
		TreeNode* n = tree->root;
		TreeNode* p = NULL;
		int (*comparator)(void*,void*) = tree->cmpfunc;
		int cmpresult;
		/* Find correct place for the new node inside while loop */
		while (n) {
			p = n;
			cmpresult = comparator(x,n->x);
			if (cmpresult>0) {
				n = n->right;
			}
			else if (cmpresult<0){
				n = n->left;
			}
			else {
				return n;
			}
		}
		return BTreeAddActual(tree, x, p);
	}
	return NULL;
}

/**
 * BTreeReplace
 *	Adds an element into BTree structure performing tree balancing
 *	using AVL methods when necessary. If element exists, replaces.
 * 
 * Parameters:
 *	tree  the tree into which the element should be added to
 *	x     the pointer which is saved into the tree
 * 
 * Returns:
 *	the previous element or NULL if added new or failed
 */
DLLEXPORT(void*) BTreeReplace(BTree *tree, void *x) {
	void *old = NULL;
	if (tree) {
		/* Initialization */
		TreeNode* n = tree->root;
		TreeNode* p = NULL;
		int (*comparator)(void*,void*) = tree->cmpfunc;
		int cmpresult;
		/* Find correct place for the new node inside while loop */
		while (n) {
			p = n;
			cmpresult = comparator(x,n->x);
			if (cmpresult>0) {
				n = n->right;
			}
			else if (cmpresult<0){
				n = n->left;
			}
			else {
				old = n->x;
				BTreeDel(tree, n);
			}
		}
		BTreeAddActual(tree, x, p);
	}
	return old;
}

/**
 * BTreeAddActual
 *	Adds an element into BTree structure performing tree balancing
 *	using AVL methods when necessary. Used by BTreeAdd and Replace.
 * 
 * Parameters:
 *	tree	the tree into which the element should be added to
 *	x		the pointer which is saved into the tree
 *  p       the parent node under which the new element is added
 * 
 * Returns:
 *	the TreeNode of added element
 *  NULL if the element exists already
 */
TreeNode *BTreeAddActual(BTree *tree, void *x, TreeNode *p) {
	/* Create the new treenode and assign it default values */
	int (*comparator)(void*,void*) = tree->cmpfunc;
	int lheight, rheight, mheight;
	TreeNode *n = (TreeNode*)malloc(sizeof(TreeNode));
	if (!n) return NULL;
	n->x=x;
	n->parent=p;
	n->left=NULL;
	n->right=NULL;
	n->height=1;
	if (p) {
		if (comparator(x,p->x)>0) {
			p->right=n;
		}else{
			p->left=n;
		}
	}
	else {
		tree->root = n;
	}
	/* Count the new tree heights and perform AVL adjustment if necessary */
	while (p) {
		lheight = 0;
		rheight = 0;
		if (p->left) lheight=p->left->height;
		if (p->right) rheight=p->right->height;
		mheight = (lheight>rheight)?lheight:rheight;
		if (p->height != ++mheight) {
			/* the height of this node has been changed, adjust */
			p->height = mheight;
			p = AVLTreeCheckNode(p,tree);
		}
		else break; /* if this node has previous length,
			necessarily all its parents have also */
		p = p->parent;
	}
	return n;
}

/**
 * BTreeDel
 *	Deletes a node from AVL tree.
 * 
 * Parameters:
 *	tree	the tree from which the node is to be deleted
 *	n		the pointer to the TreeNode to be removed
 * 
 * Returns:
 *	0 if the operation fails (n == NULL)
 *	1 if the operation succeeds
 *	2 if the tree became empty
 * 
 * Notes:
 *	Undefined behavior if trying to delete from a tree in which
 *	the node does not exist!
 */
DLLEXPORT(int) BTreeDel(BTree *tree, TreeNode *n) {
	int lheight, rheight, mheight;
	TreeNode *n1, *p1, *r, *p;
	if (!n || !tree) return 0;
	r = tree->root;
	p = n->parent;
	if (n->height==1) {
		/* Special case: Leaf node */
		if (n==r) {
			/* Very special case: Root+Leaf node = The only node */
			tree->root = NULL;
			free(n);
			return 2;
		}
		if (p->left==n) p->left=NULL;
		else p->right=NULL;
		p1 = p;
	}
	else {
		/* For non-leaf nodes */
		int bal = AVLTreeCountBalance(n);
		if (bal<0) {
			/* Left side is longer, take replacement from there */
			n1 = n->left;
			while (n1) {
				p1=n1;
				n1=n1->right;
			}
			n1=p1;
			p1=n1->parent;
			if (p1 != n) {
				if (n1->left) {
					p1->right=n1->left;
					if (p1->right) p1->right->parent=p1;
					p1=p1->right;
				}
				else {
					p1->right = NULL;
				}
			}
		}
		else {
			/* Right side is longer or heights are equal */
			n1 = n->right;
			while (n1) {
				p1=n1;
				n1=n1->left;
			}
			n1=p1;
			p1=n1->parent;
			if (p1 != n) {
				if (n1->right) {
					p1->left=n1->right;
					if (p1->left) p1->left->parent=p1;
					p1=p1->left;
				}
				else {
					p1->left = NULL;
				}
			}
		}
		
		if (p) {
			if (p->left==n) p->left=n1;
			else p->right=n1;
		}
		else {
			tree->root = n1;
		}
		n1->parent = n->parent;
		if (n1 != n->left) {
			n1->left = n->left;
			if (n1->left) n1->left->parent = n1;
		}
		if (n1 != n->right) {
			n1->right = n->right;
			if (n1->right) n1->right->parent = n1;
		}
	}
	if (p1 == n) p1 = n1;
	free(n);	/* The new node has already been assigned */
	
	/* Make sure the tree has correct height values and is balanced */
	while (p1) {
		lheight = 0;
		rheight = 0;
		if (p1->left) lheight=p1->left->height;
		if (p1->right) rheight=p1->right->height;
		mheight = (lheight>rheight)?lheight:rheight;
		if (p1->height != ++mheight) {
			p1->height = mheight;
		}
		p1 = AVLTreeCheckNode(p1,tree);
		p1 = p1->parent;
	}
	return 1;
}

/**
 * AVLTreeCheckNode
 *	Checks selected tree node in AVL Tree and balances
 *	it if necessary
 * 
 * Parameters:
 *	A	the node to be checked
 *	T	the tree in which the node A is
 * 
 * Returns:
 *	a pointer to the node that replaced A in the logical hierarchy
 */
TreeNode* AVLTreeCheckNode(TreeNode* A, BTree* T) {
	TreeNode *B, *C, *p;
	int bal2, bal = AVLTreeCountBalance(A);
	p = A->parent;
	if (bal > 1) {
		/* Right side too long */
		B = A->right;
		if ((bal2 = AVLTreeCountBalance(B))>0) {
			/* Rotate left */
			A->right = B->left;
			if (A->right)
				A->right->parent = A;
			B->left = A;
			A->parent = B;
			A->height -= 2;
		}
		else {
		/* Double rotate left */
			C = B;
			B = C->left;
			C->left = B->right;
			if (C->left)
				C->left->parent = C;
			B->right = C;
			C->parent = B;
			A->right = B->left;
			if (A->right)
				A->right->parent = A;
			B->left = A;
			A->parent = B;
			A->height -= 2;
			B->height++;
			if (bal2<0) C->height--;
			else B->height++;
		}
	}
	else if (bal < -1) {
		/* Left side too long */
		B = A->left;
		if ((bal2 = AVLTreeCountBalance(B))<0) {
			/* Rotate right */
			A->left = B->right;
			if (A->left)
				A->left->parent = A;
			B->right = A;
			A->parent = B;
			A->height -= 2;
		}
		else {
			/* Double rotate right */
			C = B;
			B = C->right;
			C->right = B->left;
			if (C->right)
				C->right->parent = C;
			B->left = C;
			C->parent = B;
			A->left = B->right;
			if (A->left)
				A->left->parent = A;
			B->right = A;
			A->parent = B;
			A->height -= 2;
			B->height++;
			if (bal2>0)C->height--;
			else B->height++;
		}
	}
	else
	{
		/* No balancing needed */
		return A;
	}
	
	/* Make the new node correctly parented */
	B->parent=p;
	if (p==NULL) {
		T->root = B;
	}
	else {
		if (p->left==A) p->left = B;
		else p->right = B;
	}
	return B;
}

/**
 * AVLTreeCountBalance
 *	Counts the balance value of selected node
 * 
 * Parameters:
 *	n	the node to be checked
 * 
 * Returns:
 *	an integer resembling the balance, should be -2..2
 */
int AVLTreeCountBalance(TreeNode* n) {
	int lh = 0, rh = 0;
	if (!n) return 0;
	if (n->left) lh=n->left->height;
	if (n->right) rh=n->right->height;
	return rh-lh;
}

/**
 * freeBTree
 *  Frees the memory of the BTree structure
 *
 * Parameters:
 *  tree      the pointer to the BTree struct to be freed
 *  freefunc  pointer to void(void*) function to free the data stored
 *
 * Returns:
 *  nothing
 */
DLLEXPORT(void) freeBTree(BTree *tree, FREEFUNC *freefunc) {
	if (tree) {
		TreeNode* n = tree->root;
		if (n)
			freeBTreeNode(n,freefunc);
		free(tree);
	}
}

/**
 * freeBTreeNode
 *  Recursive function to free memory
 *
 * Parameters:
 *  n         pointer to TreeNode
 *  freefunc  pointer to void(void*) function to free the data stored
 *
 * Returns:
 *  nothing
 */
DLLEXPORT(void) freeBTreeNode(TreeNode *n, FREEFUNC *freefunc) {
	if (n->left) freeBTreeNode(n->left,freefunc);
	if (n->right) freeBTreeNode(n->right,freefunc);
	if (freefunc) freefunc(n->x);
	free(n);
}

/**
 * iterateTree
 *   Starts iterating a BTree structure
 * 
 * Parameters:
 *   tree  pointer to BTree
 * 
 * Returns:
 *   first TreeNode pointer or NULL if empty tree
 */
DLLEXPORT(TreeNode*) iterateTree(BTree *tree) {
	TreeNode *next, *self = tree->root;
	if (!self) return NULL;
	while ((next=self->left)!=NULL) {
		self=next;
	}
	return self;
}

/**
 * iterateTreeNext
 *   continues iteration
 * 
 * Parameters:
 *   self  current TreeNode pointer
 * 
 * Returns:
 *   next TreeNode pointer or NULL if iteration completed
 */
DLLEXPORT(TreeNode*) iterateTreeNext(TreeNode *orig) {
	TreeNode *next = orig, *self = orig;
	
	if (self->right == NULL) {
		self = self->parent;
		if (!self) return NULL;
		while (self->right == next) {
			next = self;
			self = self->parent;
			if (!self) return NULL;
		}
	}
	else {
		self = self->right;
		while ((next=self->left)!=NULL) {
			self=next;
		}
	}
	return self;
}

DLLEXPORT(void) BTree_debug_dump(BTree *tree) {
	TreeNode *iterator = iterateTree(tree);
	fprintf(stderr, "*** BTree_debug_dump ***\n");
	fprintf(stderr, "Tree @ %p\n", tree);
	while(iterator) {
		fprintf(stderr, "node @ %p   (height: %d) :\n", iterator, iterator->height);
		fprintf(stderr, "            parent: %p\n", iterator->parent);
		fprintf(stderr, "  left: %p    right: %p\n", iterator->left, iterator->right);
		fprintf(stderr, "         element: %p\n", iterator->x);
		
		iterator = iterateTreeNext(iterator);
	}
}
